#ifndef SD_ROTATION2_HPP
#define SD_ROTATION2_HPP

#include "config.hpp"


namespace MT {

// -------------------------------------------------------------------
// projection
// -------------------------------------------------------------------

//! Find a rotation matrix that is nearest to a 3x3 matrix
//! (in terms of Frobenius norm)
//! Input:
//!   A = a 3x3 matrix
//! Output:
//!   R = a 3x3 rotation matrix closed to A
//!
//! This function doesn't handle the case that A'A is singular.
template<typename _Scalar>
Eigen::Matrix<_Scalar, 3, 3> rotation_project(Eigen::Matrix<_Scalar, 3, 3> const &A)
{
    Eigen::SelfAdjointEigenSolver<Eigen::Matrix<_Scalar, 3, 3> > eigensolver(A.transpose()*A); // eig(A' A)
    Eigen::Matrix<_Scalar, 3, 1> D = eigensolver.eigenvalues(); // A' A = V D V'
    Eigen::Matrix<_Scalar, 3, 3> const &V = eigensolver.eigenvectors();

    _Scalar det = D[0]*D[1]*D[2];
    ei_assert(det != 0 && "This function doesn't handle the case that A'A is singular.");

	// fix the diagonal values
	int i; D.minCoeff(&i);
    D[0] = 1 / sqrt(D[0]);
    D[1] = 1 / sqrt(D[1]);
    D[2] = 1 / sqrt(D[2]);
	if(det < 0) D[i] = -D[i];

    return A*V*(D.asDiagonal())*(V.transpose()); // A V D V'
}


// -------------------------------------------------------------------
// norm and distance
// -------------------------------------------------------------------

//! Compute ||R(q1)-R(q2)||, where R(x) is the rotation matrix of quaternion x, assuming that
//! both 'q1' and 'q2' are normalised
//! normType:
//!   0 = L0-norm
//!   1 = L1-norm
//!   2 = L2-norm
//!   3 = squared L2-norm
//!   4 = L_\infty-norm
template<typename _Scalar>
_Scalar rotation_norm(Eigen::Quaternion<_Scalar> const &q1, Eigen::Quaternion<_Scalar> const &q2, int normType=2)
{
    ei_assert((0 <= normType && normType <= 4) && "normType must be in [0, 4]");
    Eigen::Matrix<_Scalar, 3, 3> r;
    r.noalias() = (q1.inverse()*q2).normalized().toRotationMatrix();
    r(0,0) -= 1; r(1,1) -= 1; r(2,2) -= 1;
    switch(normType)
    {
    case 0: return r.minCoeff();
    case 1: return r.cwiseAbs().sum();
    case 2: return r.norm();
    case 3: return r.squaredNorm();
    }
    return r.maxCoeff();
}


//! Compute the distance of two rotations measured in the space of
//! normalised quaternions. The distance here is defined as the geodesic
//! distance of two quaternions that represent the given rotations. Note
//! that since quaternions q and -q represent the same rotation, we need
//! to measure the geodesic distance from q1 to q2 and that from -q1 to q2,
//! and take the minimum distance.
//!
//! The geodesic distance is the angle (in rad) between the two quaternions
//! Solution:
//!   a1 = acos(dot(q1,q2)); angle = min(a1, pi-a1); return angle
//! thus 0 <= angle <= M_PI/2
template<typename _Scalar>
_Scalar quaternion_dist(Eigen::Quaternion<_Scalar> const &q1, Eigen::Quaternion<_Scalar> const &q2)
{
    _Scalar r = q1.dot(q2);
	if(r < -1.0 || r > 1.0) return 0;
    r = acos(r); return r <= M_PI_2? r: M_PI-r;
}




// -------------------------------------------------------------------
// mean
// -------------------------------------------------------------------


//! Returns the 'rotation mean' of a weighted array of rotations
//! Input:
//!   rotations = list of 3x3 rotation matrices (each matrix is flattened as a 9x1 vector)
//!   weights = sparse weight distribution
//!   method : int -- defines what 'rotation mean' means
//!     0 = Euclidean rotation mean, i.e. arithemtic mean projected onto the rotation space
//!     1 = Riemannian rotation mean, not yet supported
//! Output:
//!   rotation mean, as a 3x3 rotation matrix (not flattened)
template<typename _Scalar>
Eigen::Matrix<_Scalar, 3, 3> rotation_mean(Eigen::Matrix<_Scalar, 9, Eigen::Dynamic> const &rotations,
    Eigen::SparseVector<_Scalar> const &weights, int method=0)
{
    ei_assert(method == 0 && "Only the Euclidean mean is supported at the moment.");

    // compute the arithmetic mean
    Eigen::Matrix<_Scalar, 9, 1> r;
    r.noalias() = (rotations*weights) / weights.sum();

    // project it to the rotation space and return the quaternion representation
    return rotation_project(r.conservativeResize(3,3));
}



//! Returns the 'rotation mean' of a weighted array of rotations
//! Input:
//!   rotations = list of normalised quaternions
//!   weights = sparse weight distribution
//!   method : int -- defines what 'rotation mean' means
//!     0 = Euclidean rotation mean, i.e. arithemtic mean projected onto the rotation space
//!     1 = Riemannian rotation mean, not yet supported
//! Output:
//!   rotation mean, as a normalised quaternion
//! Remark:
//!   This approach actually wraps around rotation_mean().
template<typename _Scalar>
Eigen::Quaternion<_Scalar> quaternion_mean(Eigen::Matrix<_Scalar, 4, Eigen::Dynamic> const &rotations,
    Eigen::SparseVector<_Scalar> const &weights, int method=0)
{
    ei_assert(method == 0 && "Only the Euclidean mean is supported at the moment.");

    // compute the arithmetic mean
    Eigen::Matrix<_Scalar, 3, 3> r = Eigen::Matrix<_Scalar, 3, 3>::Zero();
    for (typename Eigen::SparseVector<_Scalar>::InnerIterator it(weights); it; ++it)
        r.noalias() += it.value() *
            Eigen::Map<Eigen::Quaternion<_Scalar> >(rotations.col(it.index()).data()).toRotationMatrix();
    r.noalias() /= weights.sum();

    // project it to the rotation space and return the quaternion representation
    return Eigen::Quaternion<_Scalar>(rotation_project(r)).normalized();
}







// -------------------------------------------------------------------
// medoid
// -------------------------------------------------------------------

//! Returns the 'rotation medoid' of a weighted array of rotations
//! Input:
//!   rotations = list of normalised quaternions
//!   weights = sparse weight distribution
//! Output:
//!   rotation medoid, as a normalised quaternion
//! Remark:
//!   This approach uses quaternion_dist() to measure the distance between two quaternions.
template<typename _Scalar>
Eigen::Quaternion<_Scalar> quaternion_medoid(Eigen::Matrix<_Scalar, 4, Eigen::Dynamic> const &rotations,
    Eigen::SparseVector<_Scalar> const &weights, int method=0)
{
    int bi = -1, i, j, N = weights.nonZeros();

    // return identity quaternion if there's no vote
    if(!N) return Eigen::Quaternion<_Scalar>();

    // create a temporary array for fast processing
    std::vector<Eigen::Quaternion<_Scalar>, Eigen::aligned_allocator<Eigen::Quaternion<_Scalar> > > R(N);
    i = 0;
    for(typename Eigen::SparseVector<_Scalar>::InnerIterator it(weights); it; ++it)
        R[i++].coeffs() = rotations.col(it.index());

    // TODO: sort the weights instead

    real w, bw=1e100;
    for(i = 0; i < N; ++i)
    {
        for(w = 0, j = 0; j < N && w < bw; ++j)
            w += weights._valuePtr()[j] * quaternion_dist(R[i], R[j]);
        if(w < bw) { bi = i; bw = w; }
    }
    return R[bi];
}




//! Returns the 'rotation medoid' of a weighted array of rotations
//! Input:
//!   rotations = list of 3x3 rotation matrices (each matrix is flattened as a 9x1 vector)
//!   weights = sparse weight distribution
//!   normType: see rotation_norm()
//! Output:
//!   rotation medoid, as a 3x3 rotation matrix (not flattened)
//! Remark:
//!   This approach uses rotation_norm() to measure the distance between two rotation matrices.
template<typename _Scalar>
Eigen::Matrix<_Scalar, 3, 3> rotation_medoid(Eigen::Matrix<_Scalar, 9, Eigen::Dynamic> const &rotations,
    Eigen::SparseVector<_Scalar> const &weights, int normType=2)
{
    // this method is a bit slow but it could be improved by e.g. using Euclidean rotation mean as an initial guess

    typename Eigen::SparseVector<_Scalar>::InnerIterator it;
    int bi = -1, i, j, N = weights.nonZeros();

    // return identity rotation matrix if there's no vote
    if(!N) return Eigen::Matrix<_Scalar, 3, 3>::Identity();

    // create a temporary array for fast processing
    std::vector<Eigen::Map<Eigen::Matrix<_Scalar, 3, 3> > > R(N);
    i = 0;
    for(typename Eigen::SparseVector<_Scalar>::InnerIterator it(weights); it; ++it)
        R[i++] = rotations.col(it.index()).data();

    // TODO: sort the weights instead

    real w, bw = 1.0e100;
    for(i = 0; i < N; ++i)
    {
        for(w = 0, j = 0; j < N && w < bw; ++j)
            w += weights._valuePtr()[j] * rotation_norm(R[i], R[j], normType);
        if(w < bw) { bi = i; bw = w; }
    }

    return R[bi];
}


} // namespace MT

#endif
